package constructdata.TrainDataConstructors;

import codetree.CodeTree;
import codetree.CodeTreeOperation;
import codetree.TreeNode;
import treeview.DisplayTreeView;
import treeview.TreeView;

import java.io.*;
import java.util.LinkedList;
import java.util.List;

/**
 * Construct training data (Preceding Context)
 * Created by wangxin on 2017/4/8.
 */
public class ConstructPreTrainData {

    private boolean isDebug = false;

    // The marks
    private final String[] filterSigns = {"[","]"};

    // The construct trees
    private LinkedList<CodeTree> trees;
    // The predication holes correspondent to the constructTrees;
    private LinkedList<String> predications;
    // CodeTreeOperator
    private CodeTreeOperation operator;
    // The predicate node's complete class name
    private LinkedList<String> classnames;
    // The node who is the parent of the predicated one
    private LinkedList<TreeNode> parents;
    // The size of hole
    private  LinkedList<String> holesizes;
    // The block predictions
    private LinkedList<String> blockpredictions;
    private String blockprediction;
    // The original statements
    private LinkedList<String> originalStatements;
    // The variable names
    private LinkedList<String> variableNames;
    private final String[] holeParentAdjust = {"else","elseif","catch","finally","case","default"};

    // test
    private int count = 0;
    /**
     * Constructor
     * */
    public ConstructPreTrainData(){
        operator = new CodeTreeOperation();
    }

    /**
     * Construct training tree from the input code tree
     * Return the LinkedList with 2 list:
     * 1. The training trees with some continuous holes;
     * 2. The predication holes correspondent to the training trees;
     * 3. And save the trees and predicates to the file
     * */
    public void construct(CodeTree tree,
                          FileWriter treeWriter,FileWriter predictionWriter,FileWriter classWriter, FileWriter generationNodeWriter,FileWriter treeSentenceWriter, FileWriter jarWriter, FileWriter holeSizeWriter,
                          FileWriter traceWriter, // trace back
                          FileWriter blockpredictionWriter, // block of predictions (more lines)
                          FileWriter originalStatementsWriter, // original statements
                          FileWriter variableNamesWriter,// variable names
                          FileWriter linesWriter,// lines writer
                          boolean isCompleteFlag){
        LinkedList<LinkedList> result = getConstructTrainingData(tree);
        List<CodeTree> treeList = result.get(0);
        List<String> predictionList = result.get(1);
        List<String> classList = result.get(2);
        List<TreeNode> generationNodeList = result.get(3);
        List<String> holeSizeList = result.get(4);// record size of hole
        List<String> blockpredictionList = result.get(5);
        List<String> originalStatementsList = result.get(6);
        List<String> variableNameList = result.get(7);

        //FindJarHandler findJarHandler = new FindJarHandler();
        for (int i = 0; i < treeList.size(); i++) {
            try {
                //String jar = findJarHandler.getPackage(classList.get(i));
                //if(jar != null) {
                CodeTree tempTree = treeList.get(i);

                operator.saveRegularizedTreeInFile(operator.regularization(tempTree), treeWriter);
                operator.saveTrainingPredictionInFile(predictionList.get(i), predictionWriter);
                operator.saveTrainingPredictionInFile(classList.get(i), classWriter);
                //operator.saveTrainingPredictionInFile(jar, jarWriter);
                int parentnum = generationNodeList.get(i).getSerialNumber();
                operator.saveTrainingPredictionInFile(parentnum + " " + ((parentnum != 0) ? generationNodeList.get(i).getCompleteMethodDeclaration() : ""), generationNodeWriter);
                operator.saveTreeStringFormatInFile(tempTree, treeSentenceWriter, isCompleteFlag);

                operator.saveTrainingPredictionInFile(holeSizeList.get(i), holeSizeWriter);

                operator.saveTrainingPredictionInFile(tree.getFunctionTrace(), traceWriter);
                operator.saveTrainingPredictionInFile(blockpredictionList.get(i), blockpredictionWriter);

                operator.saveTrainingPredictionInFile(originalStatementsList.get(i), originalStatementsWriter);
                operator.saveTrainingPredictionInFile(variableNameList.get(i), variableNamesWriter);

                operator.saveTrainingPredictionInFile((tempTree.getLines(tempTree.getRoot()) + 1)+"",linesWriter);
                //}
            } catch (Exception e) {
                System.err.println(e.getMessage());
            } catch (Error e){
                System.err.println(e.getMessage());
            }
        }
    }

    /**
     * Construct training tree from the input code tree
     * Return the LinkedList with 2 list:
     * 1. The training trees with some continuous holes;
     * 2. The predication holes correspondent to the training trees;
     * */
    public LinkedList<LinkedList> getConstructTrainingData(CodeTree completeTree){
        // Init
        trees = new LinkedList<>();
        predications = new LinkedList<>();
        classnames = new LinkedList<>();
        parents = new LinkedList<>();
        holesizes = new LinkedList<>();
        blockpredictions = new LinkedList<>();
        originalStatements = new LinkedList<>();
        variableNames = new LinkedList<>();

        // Construct
        int count = completeTree.getTotalNumber();
        for (int i = 2; i <= count; i++) {// construct from the second node
            construct(completeTree, i);
        }

        // Return
        LinkedList<LinkedList> result = new LinkedList<>();
        result.addLast(trees);
        result.addLast(predications);
        result.addLast(classnames);
        result.addLast(parents);
        result.addLast(holesizes);
        result.addLast(blockpredictions);
        result.addLast(originalStatements);
        result.addLast(variableNames);

        return result;
    }

    /**
     * @param completeTree: training data from this tree
     * @param serialNumber: the first node to be remove is with this serial number
     * */
    private void construct(CodeTree completeTree, int serialNumber){
        // Make the replica of the completeTree to avoid destroying.
        CodeTree tree = copyCodeTree(completeTree);
        operator.setSerialNumberofEachNode(tree);

        // Initialize the blockprediction
        blockprediction  = "";

        if(tree.getTotalNumber() >= serialNumber){
            boolean isSuccess = true;
            List predicts = null;
            if(tree.getTotalNumber() == 1){// only root node.
                isSuccess = false;
            }
            else if((predicts = remove(tree,serialNumber)) == null
                    ){
                isSuccess = false;
            }
            if(isSuccess){
                count++;
                if(isDebug) {
                    System.out.println(count + ". " + predicts);
                }
                predications.addLast((String)predicts.get(0));
                trees.addLast(tree);
                classnames.addLast((String)predicts.get(1));
                TreeNode generationNode = (TreeNode)predicts.get(2);
                parents.addLast(generationNode);

                // Add hole node
                addHole(tree, generationNode);

                // set as 1.
                holesizes.addLast("1");
                blockpredictions.addLast(blockprediction.trim());
                if(isDebug) {
                    System.out.println(blockprediction);
                    displayTree(tree, true, blockprediction);
                }

                // statements
                originalStatements.addLast((String)predicts.get(3));
                // variable names
                variableNames.addLast((String) predicts.get(4));

            }
        }
    }

    private void addHole(CodeTree tree, TreeNode generationNode) {
        TreeNode hole = new TreeNode();
        hole.setClassName("hole");
        hole.setCompleteClassName("hole");
        hole.setMethodName("");
        hole.setCompleteMethodName("");
        hole.setAddMethodName(false);
        tree.addNode(tree.getTreeNode(generationNode.getSerialNumber()), hole);
        operator.setSerialNumberofEachNode(tree);
    }

    /**
     * Remove the node with serialNumber in the code tree
     * Remove it, including its children(except the remain node, e.g control)
     * Find the first ancestor node that under control node, remove the next brothers of this ancestor
     * Return its complete method declaration
     * */
    private LinkedList<Object> remove(CodeTree tree, int serialNumber){
        LinkedList<Object> result = new LinkedList<>();

        TreeNode node =  tree.getTreeNode(serialNumber);
        if(node == null){// Not exist
            return null;
        }
        TreeNode parent = node.getParentNode();
        String label = node.getCompleteMethodDeclaration();
        String className = node.getCompleteClassName();

        String statement = node.getStatement();
        List<String> previousVariableNames = node.getPreviousVariableNames();
        String variablename = previousVariableNames.size()>0?previousVariableNames.get(0):"";
        for (int i = 1; i < previousVariableNames.size(); i++) {
            variablename += " " + previousVariableNames.get(i);
        }

        // record predictions of this node
        blockprediction += " " + node2String(node);

        // Find the first ancestor node that under control node
        TreeNode ancestor = findAncestorUnderControl(node);

        // Remove the next brother of this ancestor
        TreeNode elderAncestor;
        while(ancestor != null) {
            elderAncestor = ancestor.getParentNode();
            if (elderAncestor != null) {
                int index = elderAncestor.getChildNodes().indexOf(ancestor);
                int nextbrothernum = elderAncestor.getChildNodes().size() - (index + 1);
                for (int i = 1; i <= nextbrothernum; i++) {
                    blockprediction += " " + node2String(elderAncestor.getChildNodes().get(index + 1));
                    elderAncestor.getChildNodes().remove(index + 1);
                }
            }
            ancestor = findAncestorUnderControl(elderAncestor);
        }

        // Remove the node
        if (parent != null) {
            parent.getChildNodes().remove(node);
        }

        // Reorder nodes
        operator.setSerialNumberofEachNode(tree);
        result.addLast(label);
        result.addLast(filterSigns(className));

        // Reassign hole parent
        if(isReassignHoleParent(label)){
            parent = node.getHoleParentNode();
        }

        if(parent == null) {
            parent = new TreeNode();
            parent.setSerialNumber(0);// the root do not have a parent
        }
        result.addLast(parent);
        result.addLast(statement);
        result.addLast(variablename);

        return result;
    }

    private TreeNode findAncestorUnderControl(TreeNode node) {
        TreeNode parent = null;
        TreeNode tmpnode = node;
        while((parent = tmpnode.getParentNode())!= null){
            if(parent.isControl()){
                return tmpnode;
            }
            tmpnode = parent;
        }
        return null;
    }

    // DFS node to string
    private String node2String(TreeNode treeNode) {
        String buff = treeNode.getCompleteMethodDeclaration();
        for (int i = 0; i < treeNode.getChildNodes().size(); i++) {
            buff += " " + node2String(treeNode.getChildNodes().get(i));
        }
        return buff;
    }

    // for testing ...
    private void displayTree(CodeTree codeTree, boolean isCompleteFlag,String title) {
        TreeView treeView = new TreeView();
        treeView.convertCodeTree(codeTree, isCompleteFlag);
        new DisplayTreeView(treeView.getTree(),title);
    }

    // filter the signs
    private String filterSigns(String string){
        for (String sign: filterSigns) {
            string = string.replace(sign, "");
        }
        return string;
    }

    // Reassign hole parent.
    private boolean isReassignHoleParent(String label){
        for(String l : holeParentAdjust){
            if(l.equals(label)){
                return true;
            }
        }
        return false;
    }

    private CodeTree copyCodeTree(CodeTree completeTree) {
        try {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new ObjectOutputStream(bos);
            oos.writeObject(completeTree);
            ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray());
            ObjectInputStream ois = new ObjectInputStream(bis);
            CodeTree tree = (CodeTree) ois.readObject();
            return tree;
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
        return null;
    }

}
